# 版本: 1.0.0
# 作者: 陈振玺
# 功能: CSV矩阵统计分析工具

import pandas as pd
import numpy as np
import argparse
from pathlib import Path

class StatsCalculator:
    def __init__(self, input_file, value_column, output_file, cluster_columns=None):
        self.input_file = Path(input_file)
        self.value_column = value_column
        self.output_file = Path(output_file)
        self.cluster_columns = cluster_columns.split() if cluster_columns else None
        self.data = None
    
    def load_data(self):
        """加载CSV数据"""
        try:
            if self.input_file.suffix.lower() == '.csv':
                self.data = pd.read_csv(self.input_file)
            else:
                print(f"错误: 不支持的文件格式 {self.input_file.suffix}")
                return False
            
            print(f"成功加载数据: {len(self.data)} 行, {len(self.data.columns)} 列")
            return True
            
        except Exception as e:
            print(f"加载数据失败: {e}")
            return False
    
    def validate_columns(self):
        """验证列名是否存在"""
        # 验证数值列
        if self.value_column not in self.data.columns:
            print(f"错误: 数值列 '{self.value_column}' 不存在")
            print(f"可用列名: {list(self.data.columns)}")
            return False
        
        # 检查数值列是否为数值型
        if not pd.api.types.is_numeric_dtype(self.data[self.value_column]):
            print(f"错误: 数值列 '{self.value_column}' 不是数值型")
            return False
        
        # 验证聚类列
        if self.cluster_columns:
            for col in self.cluster_columns:
                if col not in self.data.columns:
                    print(f"错误: 聚类列 '{col}' 不存在")
                    print(f"可用列名: {list(self.data.columns)}")
                    return False
        
        return True
    
    def calculate_stats(self):
        """计算统计量"""
        try:
            # 移除数值列中的空值
            clean_data = self.data.dropna(subset=[self.value_column])
            
            if len(clean_data) == 0:
                print("警告: 数值列中没有有效数据")
                return pd.DataFrame()
            
            if self.cluster_columns:
                # 按聚类列分组统计
                grouped = clean_data.groupby(self.cluster_columns)[self.value_column]
                
                # 计算各种统计量
                stats_df = pd.DataFrame({
                    'Count': grouped.count(),
                    'Sum': grouped.sum(),
                    'Mean': grouped.mean(),
                    'Median': grouped.median(),
                    'Min': grouped.min(),
                    'Max': grouped.max(),
                    'Std': grouped.std(),
                    'Var': grouped.var()
                }).reset_index()
                
                # 计算变异系数 (CV = std/mean)
                stats_df['CV'] = stats_df['Std'] / stats_df['Mean']
                
                # 计算聚类元素个数 (每个聚类组合中包含的元素数量)
                stats_df['Cluster_Size'] = len(stats_df)
                
                # 重新排列列的顺序 - 聚类列分列输出
                cols = self.cluster_columns + ['Cluster_Size', 'Count', 'Sum', 'Mean', 'Median', 'Min', 'Max', 'Std', 'Var', 'CV']
                stats_df = stats_df[cols]
                
                print(f"聚类统计完成: 使用聚类列 {self.cluster_columns}")
                print(f"发现 {len(stats_df)} 个不同的聚类组合")
            else:
                # 全局统计
                values = clean_data[self.value_column]
                
                stats_df = pd.DataFrame({
                    'Cluster_Size': [1],  # 全局统计只有一个聚类组
                    'Count': [len(values)],
                    'Sum': [values.sum()],
                    'Mean': [values.mean()],
                    'Median': [values.median()],
                    'Min': [values.min()],
                    'Max': [values.max()],
                    'Std': [values.std()],
                    'Var': [values.var()],
                    'CV': [values.std() / values.mean()]
                })
                
                print(f"全局统计完成: 共 {len(values)} 个有效数值")
            
            return stats_df
            
        except Exception as e:
            print(f"统计计算失败: {e}")
            return None
    
    def save_results(self, stats_df):
        """保存统计结果"""
        try:
            if self.output_file.suffix.lower() == '.csv':
                stats_df.to_csv(self.output_file, index=False, encoding='utf-8')
            else:
                print(f"错误: 不支持的输出格式 {self.output_file.suffix}")
                return False
            
            print(f"结果已保存到: {self.output_file}")
            return True
            
        except Exception as e:
            print(f"保存结果失败: {e}")
            return False
    
    def process(self):
        """执行完整的统计分析流程"""
        print("=" * 50)
        print("📊 CSV矩阵统计分析工具")
        print("=" * 50)
        
        # 加载数据
        if not self.load_data():
            return False
        
        # 验证列名
        if not self.validate_columns():
            return False
        
        # 计算统计量
        stats_df = self.calculate_stats()
        if stats_df is None or len(stats_df) == 0:
            return False
        
        # 保存结果
        if not self.save_results(stats_df):
            return False
        
        # 显示摘要信息
        print("\n📋 统计摘要:")
        print(f"📁 输入文件: {self.input_file}")
        print(f"📊 数值列名: {self.value_column}")
        if self.cluster_columns:
            print(f"🔗 聚类列名: {self.cluster_columns}")
        print(f"📄 输出文件: {self.output_file}")
        print(f"📈 统计结果: {len(stats_df)} 个统计组")
        
        return True

def main():
    parser = argparse.ArgumentParser(
        description='CSV矩阵统计分析工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 全局统计
  stats-calculator input.csv score -o output.csv
  
  # 按单个列聚类统计
  stats-calculator input.csv score -o output.csv --cluster "group"
  
  # 按多个列聚类统计
  stats-calculator input.csv score -o output.csv --cluster "group category"

输出统计量包括: 聚类大小、计数、和、均值、中位数、最小值、最大值、标准差、方差、变异系数(CV)
        """
    )
    parser.add_argument('input_file', help='输入CSV文件路径')
    parser.add_argument('value_column', help='要统计的数值型列名')
    parser.add_argument('-o', '--output', required=True, help='输出CSV文件路径')
    parser.add_argument('--cluster', help='聚类列名，多个列名用空格分隔')
    
    args = parser.parse_args()
    
    calculator = StatsCalculator(
        args.input_file,
        args.value_column,
        args.output,
        args.cluster
    )
    success = calculator.process()
    
    if success:
        print("\n🎉 CSV矩阵统计分析工具执行成功!")
    else:
        print("\n❌ CSV矩阵统计分析工具执行失败!")
        exit(1)

if __name__ == '__main__':
    main()